import torch
import torchvision.transforms as transforms
from PIL import Image
import torch.nn.functional as F
from network import LeNet


def main():
    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    transform = transforms.Compose([transforms.Resize((32, 32)),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    net = LeNet()
    net.load_state_dict(torch.load('../../models/LeNet.pth'))

    im = Image.open('plane.jpg')   # 网上找了一张飞机，压缩成32*32像素
    # print(np.array(im).shape)
    im = transform(im)  # [C, H, W]
    # print(np.array(im).shape)
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]
    # print(np.array(im).shape)
    
    net.eval()
    with torch.no_grad():
        outputs = net(im)
        print(outputs.shape, F.softmax(outputs, dim=1))
        predict = torch.max(outputs, dim=1)[1].numpy()
        print(classes[int(predict)])


if __name__ == '__main__':
    main()
